# app.py
from flask import Flask, request, jsonify
import pika
import uuid
import json
import threading

from mq import MQConnection

app = Flask(__name__)
responses = {}


class RpcClient:
    def __init__(self):
        self.mq = MQConnection()
        self.channel = self.mq.get_channel()

        # 声明回调队列
        result = self.channel.queue_declare(queue='', exclusive=True)
        self.callback_queue = result.method.queue

        self.channel.basic_consume(
            queue=self.callback_queue,
            on_message_callback=self.on_response,
            auto_ack=True)

    def on_response(self, ch, method, props, body):
        if props.correlation_id in responses:
            responses[props.correlation_id] = json.loads(body)

    def call(self, data):
        corr_id = str(uuid.uuid4())
        responses[corr_id] = None

        self.channel.basic_publish(
            exchange='',
            routing_key='rpc_queue',
            properties=pika.BasicProperties(
                reply_to=self.callback_queue,
                correlation_id=corr_id,
                delivery_mode=2
            ),
            body=json.dumps(data))

        # 等待响应（最大30秒）
        timeout = 300
        while timeout > 0 and not responses[corr_id]:
            self.connection.process_data_events(time_limit=1)
            timeout -= 1

        return responses.pop(corr_id, None)


@app.route('/v1/chat/completions', methods=['POST'])
def chat_completions():
    # 验证请求格式
    if not request.json or 'messages' not in request.json:
        return jsonify({"error": "Invalid request format"}), 400

    # 构造OpenAI兼容请求
    request_data = {
        "id": f"req_{uuid.uuid4().hex[:8]}",
        "model": request.json.get('model', 'default-model'),
        "messages": request.json['messages']
    }

    # 发送RPC请求
    try:
        rpc = RpcClient()
        response = rpc.call(request_data)

        if not response:
            return jsonify({"error": "Request timeout"}), 504

        return jsonify(response)

    except pika.exceptions.AMQPConnectionError:
        return jsonify({"error": "Message queue unavailable"}), 503


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)
